#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
SAR策略回测分析脚本

指定某个时间点，分析该时间点往前20根K线的SAR策略判断详情
包括sar_entry1、sar_entry2等指标值和信号判断
"""

import os
import sys
import logging
import pandas as pd
import numpy as np
from decimal import Decimal
from datetime import datetime, timedelta
import argparse
from sqlalchemy import select
from tabulate import tabulate

# 添加项目根目录到sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入项目模块
from models.base import SessionLocal, engine, Base
from models.kline_5min import Kline5Min
from services.strategies.implementations.sar_strategy import SARStrategy
from services.exchange.interface import ExchangeType
from services.contract_config_adapter import get_contract_info, get_all_symbols, get_mark_price_round

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='SAR策略回测分析工具')
    parser.add_argument('--symbol', type=str, default='BTC-USDT', 
                        help='交易对符号，例如BTC-USDT')
    parser.add_argument('--datetime', type=str, 
                        help='分析的时间点，格式为YYYY-MM-DD HH:MM:SS，默认为当前时间')
    parser.add_argument('--count', type=int, default=20, 
                        help='需要分析的K线数量，默认为20')
    parser.add_argument('--use-backtest-entry', action='store_true', 
                        help='是否使用回测模式的开仓信号')
    parser.add_argument('--use-backtest-exit', action='store_true', 
                        help='是否使用回测模式的平仓信号')
    
    return parser.parse_args()

def init_services():
    """初始化服务"""
    # 确保数据库表存在
    Base.metadata.create_all(engine)
    
    # 检查合约配置是否可用
    logger.info("正在检查合约配置信息...")
    symbols = get_all_symbols(ExchangeType.GATE)
    if symbols:
        logger.info(f"合约配置加载成功，共有 {len(symbols)} 个合约")
    else:
        logger.warning("合约配置加载失败或无可用合约")

def get_klines(symbol: str, end_datetime: datetime, count: int = 20) -> pd.DataFrame:
    """
    获取指定时间点之前的K线数据
    
    Args:
        symbol: 交易对符号
        end_datetime: 结束时间点
        count: 获取K线数量
        
    Returns:
        DataFrame: K线数据
    """
    try:
        # 获取数据库会话
        session = SessionLocal()
        
        # 构建查询
        query = (
            select(Kline5Min)
            .where(Kline5Min.symbol == symbol)
            .where(Kline5Min.datetime <= end_datetime)
            .order_by(Kline5Min.datetime.desc())
            .limit(count + 50)  # 多获取一些数据，用于计算指标
        )
        
        # 执行查询
        result = session.execute(query).scalars().all()
        
        # 转换为DataFrame
        data = []
        for kline in result:
            data.append({
                'datetime': kline.datetime,
                'open': float(kline.open) if isinstance(kline.open, Decimal) else kline.open,
                'high': float(kline.high) if isinstance(kline.high, Decimal) else kline.high,
                'low': float(kline.low) if isinstance(kline.low, Decimal) else kline.low,
                'close': float(kline.close) if isinstance(kline.close, Decimal) else kline.close,
                'volume': float(kline.volume) if isinstance(kline.volume, Decimal) else kline.volume
            })
        
        # 检查是否有数据
        if not data:
            logger.warning(f"未找到交易对 {symbol} 的K线数据")
            return pd.DataFrame()
        
        # 创建DataFrame并按时间排序
        df = pd.DataFrame(data)
        
        # 确保datetime列存在
        if 'datetime' not in df.columns:
            logger.error(f"K线数据缺少datetime列")
            return pd.DataFrame()
            
        df = df.sort_values('datetime')
        
        # 确保所有数值列是浮点型
        for col in ['open', 'high', 'low', 'close', 'volume']:
            if col in df.columns:
                df[col] = df[col].astype(float)
            
        return df
    
    except Exception as e:
        logger.error(f"获取K线数据时出错: {e}")
        raise
    
    finally:
        session.close()

def analyze_sar_strategy(df: pd.DataFrame, use_backtest_entry: bool = False, use_backtest_exit: bool = False, count: int = 20) -> None:
    """
    分析SAR策略
    
    Args:
        df: K线数据
        use_backtest_entry: 是否使用回测模式的开仓信号
        use_backtest_exit: 是否使用回测模式的平仓信号
        count: 分析的K线数量
    """
    # 初始化SAR策略
    strategy_params = {
        'use_backtest_entry_mode': use_backtest_entry,
        'use_backtest_exit_mode': use_backtest_exit,
        'sar_entry_offset_pct': 1.5
    }
    
    strategy = SARStrategy(df['symbol'].iloc[0] if 'symbol' in df else 'Unknown', strategy_params)
    
    # 准备数据（计算技术指标）
    result_df = strategy.prepare_data(df)
    
    # 只保留最后count行数据进行分析
    analysis_df = result_df.tail(count).copy()
    
    # 添加序号列
    analysis_df['序号'] = range(1, len(analysis_df) + 1)
    
    # 格式化日期时间
    analysis_df['日期时间'] = analysis_df['datetime'].dt.strftime('%Y-%m-%d %H:%M')
    
    # 选择需要展示的列
    columns = [
        '序号', '日期时间', 'open', 'high', 'low', 'close', 
        'sar_entry1', 'sar_entry2', 'sar_exit', 'sar_filter',
        'supertrend_trend', 'adx', 'rvi',
        'entry1_buy_signal', 'entry2_buy_signal', 'sar_buy', 'trend_filter', 'sar_B',
        'entry1_sell_signal', 'entry2_sell_signal', 'sar_sell', 'sar_S',
        'long_signal', 'short_signal',
        'exit_long_signal', 'exit_short_signal'
    ]
    
    # 创建一个更易读的展示DataFrame
    display_df = pd.DataFrame()
    
    # 复制基本列
    display_df['序号'] = analysis_df['序号']
    display_df['日期时间'] = analysis_df['日期时间']
    
    # 格式化价格列
    for col in ['open', 'high', 'low', 'close']:
        display_df[col] = analysis_df[col].apply(lambda x: strategy.base_format_price(x))
    
    # 格式化SAR指标列
    for col in ['sar_entry1', 'sar_entry2', 'sar_exit', 'sar_filter']:
        display_df[col] = analysis_df[col].apply(lambda x: strategy.base_format_price(x))
    
    # 复制其他指标列
    display_df['SuperTrend'] = analysis_df['supertrend_trend'].apply(lambda x: '看多' if x == 1 else '看空')
    display_df['ADX'] = analysis_df['adx'].round(2)
    display_df['RVI'] = analysis_df['rvi'].round(2)
    
    # 格式化信号列
    for col in ['entry1_buy_signal', 'entry2_buy_signal', 'sar_buy', 'trend_filter', 'sar_B',
                'entry1_sell_signal', 'entry2_sell_signal', 'sar_sell', 'sar_S',
                'long_signal', 'short_signal', 'exit_long_signal', 'exit_short_signal']:
        display_df[col] = analysis_df[col].apply(lambda x: '✓' if x else '✗')
    
    # 重命名列，使其更易读
    column_mapping = {
        'entry1_buy_signal': 'SAR224买',
        'entry2_buy_signal': 'SAR225买',
        'sar_buy': 'SAR价格买',
        'trend_filter': '趋势过滤',
        'sar_B': 'SAR223多头',
        'entry1_sell_signal': 'SAR224卖',
        'entry2_sell_signal': 'SAR225卖',
        'sar_sell': 'SAR价格卖',
        'sar_S': 'SAR223空头',
        'long_signal': '多头信号',
        'short_signal': '空头信号',
        'exit_long_signal': '多头平仓',
        'exit_short_signal': '空头平仓'
    }
    display_df = display_df.rename(columns=column_mapping)
    
    # 计算信号总结
    signal_summary = []
    for i, row in analysis_df.iterrows():
        summary = ""
        if row['long_signal']:
            summary = "【多头开仓信号】"
        elif row['short_signal']:
            summary = "【空头开仓信号】"
        elif row['exit_long_signal']:
            summary = "【多头平仓信号】"
        elif row['exit_short_signal']:
            summary = "【空头平仓信号】"
        signal_summary.append(summary)
    
    display_df['信号总结'] = signal_summary
    
    # 打印表格
    print("\n=== SAR策略回测分析 ===")
    print(f"交易对: {df['symbol'].iloc[0] if 'symbol' in df else 'Unknown'}")
    print(f"开仓模式: {'回测' if use_backtest_entry else '实盘'}")
    print(f"平仓模式: {'回测' if use_backtest_exit else '实盘'}\n")
    
    # 将DataFrame分成两部分打印，避免表格太宽
    # 第一部分：基本价格和SAR指标
    basic_cols = ['序号', '日期时间', 'open', 'high', 'low', 'close', 
                 'sar_entry1', 'sar_entry2', 'sar_exit', 'sar_filter',
                 'SuperTrend', 'ADX', 'RVI']
    print("== 价格和SAR指标 ==")
    print(tabulate(display_df[basic_cols], headers='keys', tablefmt='pretty', showindex=False))
    
    # 第二部分：信号判断
    signal_cols = ['序号', '日期时间', 'SAR224买', 'SAR225买', 'SAR价格买', '趋势过滤', 'SAR223多头',
                  'SAR224卖', 'SAR225卖', 'SAR价格卖', 'SAR223空头',
                  '多头信号', '空头信号', '多头平仓', '空头平仓', '信号总结']
    print("\n== 信号判断 ==")
    print(tabulate(display_df[signal_cols], headers='keys', tablefmt='pretty', showindex=False))
    
    # 分析每个K线的详细情况
    print("\n== 详细分析 ==")
    for i, row in analysis_df.iterrows():
        idx = row['序号']
        dt = row['日期时间']
        
        print(f"\n[{idx}] {dt} - {'有信号' if any([row['long_signal'], row['short_signal'], row['exit_long_signal'], row['exit_short_signal']]) else '无信号'}")
        
        # 价格信息
        print(f"  价格: 开={strategy.base_format_price(row['open'])}, 高={strategy.base_format_price(row['high'])}, 低={strategy.base_format_price(row['low'])}, 收={strategy.base_format_price(row['close'])}")
        
        # SAR指标
        print(f"  SAR指标: SAR224={strategy.base_format_price(row['sar_entry1'])}, "
              f"SAR225={strategy.base_format_price(row['sar_entry2'])}, "
              f"SAR227={strategy.base_format_price(row['sar_exit'])}, "
              f"SAR223={strategy.base_format_price(row['sar_filter'])}")
        
        # 趋势指标
        print(f"  趋势指标: SuperTrend={'看多' if row['supertrend_trend'] == 1 else '看空'}, "
              f"ADX={row['adx']:.2f}, RVI={row['rvi']:.2f}")
        
        # 信号判断详细解析
        if row['long_signal']:
            print(f"  【多头开仓信号】 - 所有多头条件均满足:")
            print(f"    - SAR224买信号: {'满足' if row['entry1_buy_signal'] else '不满足'}")
            print(f"    - SAR225买信号: {'满足' if row['entry2_buy_signal'] else '不满足'}")
            print(f"    - SAR价格买条件: {'满足' if row['sar_buy'] else '不满足'}")
            print(f"    - 趋势过滤条件: {'满足' if row['trend_filter'] else '不满足'}")
            print(f"    - SAR223多头: {'满足' if row['sar_B'] else '不满足'}")
            print(f"    - SuperTrend趋势: {'看多' if row['supertrend_trend'] == 1 else '看空'}")
        
        elif row['short_signal']:
            print(f"  【空头开仓信号】 - 所有空头条件均满足:")
            print(f"    - SAR224卖信号: {'满足' if row['entry1_sell_signal'] else '不满足'}")
            print(f"    - SAR225卖信号: {'满足' if row['entry2_sell_signal'] else '不满足'}")
            print(f"    - SAR价格卖条件: {'满足' if row['sar_sell'] else '不满足'}")
            print(f"    - 趋势过滤条件: {'满足' if row['trend_filter'] else '不满足'}")
            print(f"    - SAR223空头: {'满足' if row['sar_S'] else '不满足'}")
            print(f"    - SuperTrend趋势: {'看多' if row['supertrend_trend'] == 1 else '看空'}")
        
        elif row['exit_long_signal']:
            print(f"  【多头平仓信号】- SAR227={strategy.base_format_price(row['sar_exit'])}, 收盘价={strategy.base_format_price(row['close'])}")
        
        elif row['exit_short_signal']:
            print(f"  【空头平仓信号】- SAR227={strategy.base_format_price(row['sar_exit'])}, 收盘价={strategy.base_format_price(row['close'])}")
        
        else:
            # 多头缺少哪些条件
            missing_long = []
            if not row['entry1_buy_signal']: missing_long.append("SAR224买信号")
            if not row['entry2_buy_signal']: missing_long.append("SAR225买信号")
            if not row['sar_buy']: missing_long.append("SAR价格买条件")
            if not row['trend_filter']: missing_long.append("趋势过滤")
            if not row['sar_B']: missing_long.append("SAR223多头")
            if not row['supertrend_trend'] == 1: missing_long.append("SuperTrend看多")
            
            # 空头缺少哪些条件
            missing_short = []
            if not row['entry1_sell_signal']: missing_short.append("SAR224卖信号")
            if not row['entry2_sell_signal']: missing_short.append("SAR225卖信号")
            if not row['sar_sell']: missing_short.append("SAR价格卖条件")
            if not row['trend_filter']: missing_short.append("趋势过滤")
            if not row['sar_S']: missing_short.append("SAR223空头")
            if not row['supertrend_trend'] == -1: missing_short.append("SuperTrend看空")
            
            print(f"  无交易信号 - 条件不满足:")
            print(f"    - 多头缺少条件: {', '.join(missing_long)}")
            print(f"    - 空头缺少条件: {', '.join(missing_short)}")

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 初始化服务（检查合约配置）
    init_services()
    
    # 解析时间
    if args.datetime:
        try:
            end_datetime = datetime.strptime(args.datetime, '%Y-%m-%d %H:%M:%S')
        except ValueError:
            logger.error("时间格式不正确，请使用YYYY-MM-DD HH:MM:SS格式")
            sys.exit(1)
    else:
        end_datetime = datetime.now()
    
    logger.info(f"开始分析 {args.symbol} 在 {end_datetime} 之前的 {args.count} 根K线")
    
    # 获取K线数据
    df = get_klines(args.symbol, end_datetime, args.count)
    
    if df.empty:
        logger.error("未找到符合条件的K线数据")
        sys.exit(1)
    
    # 添加symbol列
    df['symbol'] = args.symbol
    
    # 分析SAR策略
    analyze_sar_strategy(df, args.use_backtest_entry, args.use_backtest_exit, args.count)
    
    logger.info("分析完成")

if __name__ == "__main__":
    main() 